import requests
from fake_useragent import UserAgent
import json
import re
from tqdm import tqdm
import os

# URL 和请求头设置
url = "https://gfwsl.geforce.cn/services_toolkit/services/com/nvidia/services/AjaxDriverService.php?func=DriverManualLookup&psid=129&pfid=1007&osID=135&languageCode=2052&beta=null&isWHQL=0&dltype=-1&dch=1&upCRD=null&qnf=0&sort1=1&numberOfResults=10"
headers = {
    "User-Agent": UserAgent().random,
    "Referer": "https://www.nvidia.cn/geforce/drivers/",  # 设置防盗链
}

# 确保目标目录存在
output_dir = "D:/驱动/"
os.makedirs(output_dir, exist_ok=True)

# 发起请求
try:
    response = requests.get(url=url, headers=headers)
    response.raise_for_status()  # 检查请求是否成功
    data = response.json()

    # 提取下载链接
    downloadurl = [item['downloadInfo']['DownloadURL'] for item in data['IDS']]
    urls = [url for url in downloadurl if "-nsd" not in url]

    if urls:
        print("下载链接：" + urls[0])

        # 提取文件名
        name = re.findall(r"/([^/]+\.exe)", urls[0])
        if name:
            file_name = name[0]
            print("驱动名称：" + file_name)

            # 获取文件总大小
            response = requests.head(urls[0], headers=headers)
            response.raise_for_status()
            total_size = int(response.headers.get('content-length', 0))

            # 下载文件并显示进度条
            with requests.get(urls[0], headers=headers, stream=True) as r, open(os.path.join(output_dir, file_name), mode="wb") as f:
                progress_bar = tqdm(total=total_size, unit='B', unit_scale=True, unit_divisor=1024, desc=file_name)
                for chunk in r.iter_content(chunk_size=1024):
                    if chunk:  # filter out keep-alive new chunks
                        f.write(chunk)
                        progress_bar.update(len(chunk))
                progress_bar.close()

            print(f"文件已下载并保存为: {os.path.join(output_dir, file_name)}")
        else:
            print("在 URL 中未找到文件名。")
    else:
        print("未找到下载 URL。")

    wr = input("请输入是否下载 nvidia app（y下载，n或者其他不下载）：")
    if wr.lower() == "y":
        urla = "https://cn.download.nvidia.cn/nvapp/client/10.0.1.256/NVIDIA_app_beta_v10.0.1.256.exe"
        response1 = requests.get(url=urla, headers=headers)
        response1.raise_for_status()
        total_size1 = int(response1.headers.get('content-length', 0))

        with requests.get(urla, headers=headers, stream=True) as r1, open(os.path.join(output_dir, "nvidia app.exe"), mode="wb") as f1:
            progress_bar1 = tqdm(total=total_size1, unit="B", unit_scale=True, unit_divisor=1024, desc="nvidia app.exe")
            for chunk1 in r1.iter_content(chunk_size=1024):
                if chunk1:
                    f1.write(chunk1)  # 修正：应该写入 f1
                    progress_bar1.update(len(chunk1))
            progress_bar1.close()
        print("nvidia app.exe 下载成功。。。")
    else:
        print("已取消下载。。。")

except requests.RequestException as e:
    print(f"网络请求发生错误: {e}")
except Exception as e:
    print(f"发生错误: {e}")
